"""
src/odelib/assembly.py
    Create an Assembly class so we can manage collections of bodies that form one object.

Written By:
    James Thomas
    Email: jim@houseoftechnology.org
    Web: http://mission-cognition.houseoftechnology.org/

Copyright 2009-2012

This file is part of the PandaODElib distribution.

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public License
as published by the Free Software Foundation; either version 3 of
the License, or (at your option) any later version. The text of the
GNU Lesser General Public License is included with this library in
the file named LICENSE.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
"""

from pandac.PandaModules import Vec3, OdeBody

def QuatInverse(Q):
    """Panda3d inverse does not conform to any definition I've seen.  Use this instead."""
    return Q.conjugate() / Q.lengthSquared()

def QuatMultiply(Q1, Q2):
    """Panda3d multiplication order seems backwards from all the definitions I've seen. Use this instead"""
    return Q2 * Q1

def QuatFromQuaternionPair(q1, q2): 
    """Creates a quaternion q which will rotate q1 to q2 such that RotateQuat(q, q1) -> q2"""
    return QuatMultiply(q2, QuatInverse(q1))

class Assembly(object):
    def __init__(self):
        self._objDict = {}

    def _setConnectionJointInstance(self, ConnectionID, JointInstance):
        # Needed to be interchangable with basebody objects
        body = self.GetBodyFromConnectionID(ConnectionID)
        connID = ConnectionID.rsplit('.', 1)[1]
        body._setConnectionJointInstance(connID, JointInstance)

    def GetBodyFromConnectionID(self, ConnectionID):
        """Get the body associated with the specified assembly-dotted connection ID.
        Example ConnectionIDs formats are subassy1.body1.connection1, body1.connection1"""
        # Split off everything before the connection (up to the last dot)
        key = ConnectionID.rsplit('.', 1)[0]
        return self.GetObj(key)
        
    def AddObj(self, Name, Obj):
        """Add a obj to the assembly using Name for future reference.  Name can be
        any hashable type usable as a dict key. You may not use dotted sub-assembly keys here."""
        assert Name not in self._objDict  # Not ready for replaces yet
        self._objDict[Name] = Obj

    def GetNames(self):
        """Get a list of obj key values of bodies in the assembly."""
        return self._objDict.keys()

    def IsIn(self, ObjObject):
        """Return True if the obj instance is part of this assembly.  Use GetNames if
        you have the obj by key and want to know if it is in the assembly."""
        return ObjObject in self._objDict.values()

    def GetObj(self, Name):
        """Get the obj associated with Name.  Name supports the dotted sub-assembly format."""
        # Handle dotted keys
        splits = Name.split('.', 1)
        
        if len(splits) == 1:
            splits.append('')
        
        key, remainder = splits

        assert key in self._objDict
        obj = self._objDict.get(Name)
        
        if remainder:
            return obj.GetObj(remainder)
        else:
            return obj
        
    #def GetObjPosition(self, Name):
    #    """Get the obj position in world coordinates for the obj associated with Name"""
    #    return self.GetObj(Name).getPosition()

    #def GetObjQuaternion(self, Name):
    #    """Get the orientation quaternion (in world coordinates) for the obj associated with Name"""
    #    return self.GetObj(Name).getQuaternion()

    #def SetObjPosition(self, Name, Position):
    #    """Set the position of the obj associated with Name (Position is specified in world coordinates)."""
    #    self.GetObj(Name).setPosition(Position)

    #def SetObjQuaternion(self, Name, Quaternion):
    #    """Set the orientation of the obj associated with Name (Quaternion
    #    is specified in world coordinates)."""
    #    self.GetObj(Name).setQuaternion(Quaternion)

    #def GetMassAndCG(self):
    #    """Get the mass and cg of the assembly returned as mass, cg (3-tuple vector)."""
    #    sumCGMoments = (0.0, 0.0, 0.0)
    #    sumMass = 0.0
    #    for obj in self._objDict.itervalues():
    #        m, cg, I = obj.GetMassCGAndInertia()
    #        # Add the mass from this obj
    #        sumMass += m

    #        # Add the moment from this obj
    #        sumCGMoments = vsum(sumCGMoments, sprod(m, cg))

    #    return sumMass, sprod(1.0 / sumMass, sumCGMoments)

    #def MoveCGTo(self, NewPosition):
    #    """Move the entire assembly so that it's cg is located at NewPosition (3-tuple vector
    #    in world coordinates)"""
    #    m, cg = self.GetMassAndCG()

    #    # Calculate the delta between the new position and the current CG
    #    delta = vdiff(NewPosition, cg)

    #    self.MoveByDelta(delta)

    def MoveByDelta(self, Delta):
        """Move the entire assembly by the 3-tuple vector Delta (in world coordinates)."""
        # Apply the delta to all the bodies
        for obj in self._objDict.itervalues():
            obj.MoveByDelta(Delta)

    def MoveToPlaceConnectionAt(self, ConnectionID, Position):
        """Move the entire assembly such that the specified connection on the specified body
        is located at Position (3-tuple vector in world coordinates). Requires use of
        dotted connection ID notation."""
        # Get the body related to this connection
        body = self.GetBodyFromConnectionID(ConnectionID)
        connID = ConnectionID.rsplit('.', 1)[1]

        # Lookup the obj and get the location of the specified connection point.
        location = body.GetConnectionLocation(connID)

        # Calculate the delta between the new position and the connection point
        delta = Vec3(Position) - Vec3(location)

        # Move the assembly
        self.MoveByDelta(delta)

    def OrientConnectionAt(self, ConnectionID, Orientation):
        """Orient the entire assembly such that the specified connection orientation on the
        specified body is aligned to the specified 4-tuple quaternion (in world coordinates).
        Requires use of dotted connection ID notation."""

        # Get the body related to this connection
        body = self.GetBodyFromConnectionID(ConnectionID)
        connID = ConnectionID.rsplit('.', 1)[1]

        # Get the axis of the specified connection point (in world coordinates).
        connQ = body.GetConnectionOrientation(connID)

        qr = QuatFromQuaternionPair(connQ, Orientation)

        # Get the location of the specified connection point.
        location = body.GetConnectionLocation(connID)

        # Rotate the body about the connection point
        self.RotateAboutPoint(location, qr)

    def Connect(self, JointInstance, ConnectionID, Other=None, OtherConnectionID=''):
        """Make a connection using the specified joint instance between the specified connection
        on this assembly and the specified connection on the other body/assembly. If Other is not 
        specified (or None) then the connection is made to the environment."""

        # Get the body related to this connection
        body = self.GetBodyFromConnectionID(ConnectionID)
        connID = ConnectionID.rsplit('.', 1)[1]

        # Get the information for this joint location
        location = body.GetConnectionLocation(connID)
        orientation = body.GetConnectionOrientation(connID)

        if Other is not None:
            # Orient Other Body/Assembly such that the connections between the two bodies line up
            Other.MoveToPlaceConnectionAt(OtherConnectionID, location)
            Other.OrientConnectionAt(OtherConnectionID, orientation)
            
            if hasattr(Other, 'GetBodyFromConnectionID'):
                otherBody = Other.GetBodyFromConnectionID(OtherConnectionID)
            elif isinstance(Other, OdeBody):
                otherBody = Other
            else:
                raise TypeError, 'Unexpected type for Other parameter: ' + type(Other)
        else:
            otherBody = None
        
        # Make the connection.
        # If OtherBody is None then it attaches to the environment
        JointInstance.attach(body, otherBody)

        # Set the location of the joint to the prepared location
        try:
            JointInstance.setAnchor(location)
        except AttributeError:
            pass

        try:
            # Compute the joint axis based on the connection orientation
            axis = orientation.xform((0,0,1))
            JointInstance.setAxis(axis)
        except AttributeError:
            # Not all joints have a setAxis() method
            pass

        # Put the joint object into this body and the other body/assembly connection lists
        self._setConnectionJointInstance(ConnectionID, JointInstance)
        if Other is not None:
            Other._setConnectionJointInstance(OtherConnectionID, JointInstance)

    def RotateAboutPoint(self, Point, Quaternion):
        """Rotate the entire assembly about the point by the Quaternion (normalized in world coordinates)."""
        # Work on every object in the assembly.
        for obj in self._objDict.itervalues():
            obj.RotateAboutPoint(Point, Quaternion)

    def UpdateSettings(self):
        """This function must be called before each time step is executed.  It will
        call the UpdateSettings() function for every obj in the assembly."""
        for obj in self._objDict.itervalues():
            # Update the obj Settings
            obj.UpdateSettings()

